#
# Copyright (C) Since 2018 ~ 2025, Inc. All Rights Reserved
#
# @Time    : 2025/1/12-下午2:03
# @Author  : Timmy
# @Email   : tangping@cloudstream-tech.com
# @File    : terminal.py
# @Software: lottery
# @Function:
#
import asyncio

import websockets
from websockets import Data
from websockets.asyncio.client import ClientConnection


class Connection(object):
    def __init__(self):
        self.ready = asyncio.Event()
        self.connection = None

    async def connect(self):
        while True:
            try:
                self.connection: ClientConnection = await websockets.connect('ws://localhost:8000/ws/terminal')
                self.ready.set()
                break
            except ConnectionRefusedError as e:
                print(e)
                await asyncio.sleep(3)

    async def close(self):
        if self.connection:
            await self.connection.close()

    async def send_message(self, message: str):
        if not self.connection:
            await self.ready.wait()
        await self.connection.send(message)

    async def receive_message(self) -> Data | None:
        if not self.connection:
            await self.ready.wait()
        try:
            data = await self.connection.recv()
        except websockets.exceptions.ConnectionClosedOK:
            return None
        return data
